#!/usr/bin/env python
# -*- conding: utf-8 -*-

"""
@Time     : 2024/9/2 6:55
@Author   : liujingmao
@File     : 1.TCVectorDB内置Embedding使用示例.py
"""
import os

import dotenv
from langchain_community.vectorstores import TencentVectorDB
from langchain_community.vectorstores.tencentvectordb import (
    ConnectionParams,
)

dotenv.load_dotenv()

db = TencentVectorDB(
    embedding=None,
    connection_params=ConnectionParams(
        url=os.environ.get("TC_VECTOR_DB_URL"),
        username=os.environ.get("TC_VECTOR_DB_USERNAME"),
        key=os.environ.get("TC_VECTOR_DB_KEY"),
        timeout=int(os.environ.get("TC_VECTOR_DB_TIMEOUT")),
    ),
    database_name="llmops-test",
    collection_name="dataset-builtin",
)

texts = [
    "笨笨是一只很喜欢睡觉的猫咪",
    "我喜欢在夜晚听音乐，这让我感到放松。",
    "猫咪在窗台上打盹，看起来非常可爱。",
    "学习新技能是每个人都应该追求的目标。",
    "我最喜欢的食物是意大利面，尤其是番茄酱的那种。",
    "昨晚我做了一个奇怪的梦，梦见自己在太空飞行。",
    "我的手机突然关机了，让我有些焦虑。",
    "阅读是我每天都会做的事情，我觉得很充实。",
    "他们一起计划了一次周末的野餐，希望天气能好。",
    "我的狗喜欢追逐球，看起来非常开心。",
]

ids = db.add_texts(texts)
print("添加文档id列表:", ids)

"""
添加文档id列表: ['1726151741081011500--7067634663772612391-0',
 '1726151741081011500--9188772147985167064-1', 
 '1726151741081011500--3040655625764469082-2', 
 '1726151741081011500-3287730123008562965-3', 
 '1726151741081011500-2304329220341321445-4', 
 '1726151741081011500--8109298814841760202-5', 
 '1726151741081011500--867349922836396320-6', 
 '1726151741081011500-5028564518719403-7', 
 '1726151741081011500--8852105450394504360-8', 
'1726151741081011500--350854967386996021-9']

"""
# print(db.similarity_search_with_relevance_scores("我养了一只猫，叫笨笨"))
print(db.similarity_search_with_score("我养了一只猫，叫笨笨"))

"""
[(Document(page_content='笨笨是一只很喜欢睡觉的猫咪'), 0.87451), 
(Document(page_content='笨笨是一只很喜欢睡觉的猫咪'), 0.87451), 
(Document(page_content='笨笨是一只很喜欢睡觉的猫咪'), 0.87451),
 (Document(page_content='笨笨是一只很喜欢睡觉的猫咪'), 0.87451)]
"""
